'''
Author: Lee Hang
Date  : 2020/12/21

根据tf的某一层（默认最后一层）获取gradcam热力图
'''
from tensorflow.keras.models import Model
import tensorflow as tf
import numpy as np
import cv2
class GradCAM:
    def __init__(self, model, classIdx, layerName=None):
        # store the model, the class index used to measure the class
        # activation map, and the layer to be used when visualizing
        # the class activation map
        self.model = model
        self.classIdx = classIdx
        self.layerName = layerName
        # if the layer name is None, attempt to automatically find
        # the target output layer
        if self.layerName is None:
            """
            activation_5    (None, 112, 112, 96)
            activation_13   (None, 56, 56, 144)
            activation_21   (None, 28, 28, 240)
            activation_45   (None, 14, 14, 672)
            """
            self.layerName = self.find_target_layer()
            #自定义需要读取的layerName
            #self.layerName = "re_lu_196"
            #print("target layer name: " + self.layerName)

    def find_target_layer(self):
        # attempt to find the final convolutional layer in the network
        # by looping over the layers of the network in reverse order
        for layer in reversed(self.model.layers):
            # check to see if the layer has a 4D output
            #print(layer)
            if len(layer.output_shape) == 4:
                return layer.name
        # otherwise, we could not find a 4D layer so the GradCAM
        # algorithm cannot be applied
        raise ValueError("Could not find 4D layer. Cannot apply GradCAM.")

    def compute_heatmap(self, image, eps=1e-8):
        # construct our gradient model by supplying (1) the inputs
        # to our pre-trained model, (2) the output of the (presumably)
        # final 4D layer in the network, and (3) the output of the
        # softmax activations from the model
        gradModel = Model(
            inputs=[self.model.inputs],
            outputs=[self.model.get_layer(self.layerName).output,
                self.model.output])
        # record operations for automatic differentiation
        with tf.GradientTape() as tape:
            # cast the image tensor to a float-32 data type, pass the
            # image through the gradient model, and grab the loss
            # associated with the specific class index
            inputs = tf.cast(image, tf.float32)
            (convOutputs, predictions) = gradModel(inputs)
            loss = predictions[:, self.classIdx]
        # use automatic differentiation to compute the gradients
        grads = tape.gradient(loss, convOutputs)
        # compute the guided gradients
        castConvOutputs = tf.cast(convOutputs > 0, "float32")
        castGrads = tf.cast(grads > 0, "float32")
        guidedGrads = castConvOutputs * castGrads * grads
        # the convolution and guided gradients have a batch dimension
        # (which we don't need) so let's grab the volume itself and
        # discard the batch
        convOutputs = convOutputs[0]    #(7, 7, 320)
        guidedGrads = guidedGrads[0]
        # compute the average of the gradient values, and using them
        # as weights, compute the ponderation of the filters with
        # respect to the weights
        weights = tf.reduce_mean(guidedGrads, axis=(0, 1))  #(320, )
        cam = tf.reduce_sum(tf.multiply(weights, convOutputs), axis=-1)        
        # grab the spatial dimensions of the input image and resize
        # the output class activation map to match the input image
        # dimensions
        (w, h) = (image.shape[2], image.shape[1])
        camArray = cam.numpy()
        print(camArray)
        deflatAdaptor = 1e18
        heatmap = cv2.resize(camArray * deflatAdaptor, (w, h))
        # normalize the heatmap such that all values lie in the range
        # [0, 1], scale the resulting values to the range [0, 255],
        # and then convert to an unsigned 8-bit integer
        numer = heatmap - np.min(heatmap)
        denom = (heatmap.max() - heatmap.min()) + eps
        heatmap = numer / denom
        heatmap = (heatmap * 255).astype("uint8")
        # return the resulting heatmap to the calling function
        return heatmap

    def compute_heatmapplusplus(self, image, eps=1e-8):
        # 获取模型以及其输入输出
        gradModel = Model(
            inputs=[self.model.inputs],
            outputs=[self.model.get_layer(self.layerName).output,
                self.model.output])
        # with里面写公式,调用tape求导
        with tf.GradientTape() as tape:
            inputs = tf.cast(image, tf.float32)
            (convOutputs, predictions) = gradModel(inputs)
            loss = predictions[:, self.classIdx]
        # 前者对后者求导
        grads = tape.gradient(target=loss, sources=convOutputs)     #(1, 7, 7, 320)
        castConvOutputs = tf.cast(convOutputs > 0, "float32")
        castGrads = tf.cast(grads > 0, "float32")
        guidedGrads = castConvOutputs * castGrads * grads
        guidedGrads = guidedGrads[0]
        # 
        grads = np.maximum(guidedGrads, 0.)
        indicate = np.where(grads > 0, 1., 0.)      #(7, 7, 320)
        norm_factor = np.sum(grads, axis=(1, 2))
        for i in range(len(norm_factor)):
            norm_factor[i] = 1. / norm_factor[i] if norm_factor[i] > 0. else 0.
        alpha = indicate * norm_factor[:, np.newaxis, np.newaxis]   #(7, 7, 320)
        weight = np.sum(grads * alpha, axis=(0)) #(7, )
        feature = convOutputs[0]    #(7, 7, 320)
        cam = tf.multiply(feature, weight)   #(7, 1, 1)
        cam = tf.reduce_sum(cam, axis=-1)   #(7, 320)
        cam -= np.min(cam)
        cam /= np.max(cam)          #(7, 320)
        #print(cam.shape, end="")
        (w, h) = (image.shape[2], image.shape[1])
        heatmap = cv2.resize(cam.numpy(), (w, h))
        numer = heatmap - np.min(heatmap)
        denom = (heatmap.max() - heatmap.min()) + eps
        heatmap = numer / denom
        heatmap = (heatmap * 255).astype("uint8")
        # return the resulting heatmap to the calling function
        return heatmap


    def overlay_heatmap(self, heatmap, image, alpha=0.5,
            colormap=cv2.COLORMAP_VIRIDIS):
        #colormap=cv2.COLORMAP_JET):
        # apply the supplied color map to the heatmap and then
        # overlay the heatmap on the input image
        heatmap = cv2.applyColorMap(heatmap, colormap)
        output = cv2.addWeighted(image, alpha, heatmap, 1 - alpha, 0)
        # return a 2-tuple of the color mapped heatmap and the output,
        # overlaid image
        return (heatmap, output)       
